﻿using System;
using Bettzueche.RLGlue.RLComponent;
using Bettzueche.RLGlue.TaskSpec;
using Bettzueche.RLGlue.RLTypes;
using System.Diagnostics;

namespace RunMinesSarsaExperiment {
    /// <summary>
    /// This code is adapted from the Mines.cpp code written by Adam White
    /// for earlier versions of RL-Glue.
    /// See the RL-Library page: 
    /// http://library.rl-community.org/environments/mines
    /// This is a very simple discrete-state, episodic grid world that has
    /// exploding mines in it.  If the agent steps on a mine, the episode
    /// ends with a large negative reward.
    /// 
    /// The reward per step is -1, with +10 for exiting the game successfully
    /// and -100 for stepping on a mine.
    /// 
    /// This example follows my (Brian Tanner) favorite pattern of keeping the dynamics
    /// of the world fairly separate from the class that implements EnvironmentInterface.
    /// In this case, I've put it in the class WorldDescription, which is inside this
    /// same Java file as SampleMinesEnvironment.  Usually I would put it in a separate
    /// file.  This separation means that SampleMinesEnvironment doesn't need to know
    /// much about the dynamics of the world, and WorldDescription doesn't need to know
    /// much about RL-Glue.
    /// 
    /// @author Brian Tanner, Copyright 2008
    /// </summary>
    class SampleMinesEnvironment : IEnvironment {

        #region Fields

        internal const int WORLD_FREE = 0;
        internal const int WORLD_OBSTACLE = 1;
        internal const int WORLD_MINE = 2;
        internal const int WORLD_GOAL = 3;

        //WorldDescription contains the state of the world and manages the dynamics.
        WorldDescription theWorld;
        //These are used if the environment has been sent a message to use a fixed
        //starting state.
        bool fixedStartState = false;
        int startRow = 0;
        int startCol = 0;

        #endregion

        #region IEnvironment Members

        public ITaskSpec InitEnvironment() {
            //This is hard coded, but there is no reason it couldn't be automatically
            //generated or read from a file.

            int[,] world_map = {
                {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
                {1, 0, 0, 0, 0, 0, 0, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 1},
                {1, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},
                {1, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 0, 0, 0, 0, 1, 1},
                {1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 1},
                {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1}
            };


            theWorld = new WorldDescription(world_map);

            //Create a task spec programmatically.  This task spec encodes that state, action, and reward space for the problem.
            //You could forgo the task spec if your agent and environment have been created specifically to work with each other
            //ie, there is no need to share this information at run time.  You could also use your own ad-hoc task specification language,
            //or use the official one but just hard code the string instead of constructing it this way.
            
            //Specify that there will be an integer observation [0,107] for the state
            RLTypesSpec obs = new RLTypesSpec();
            obs.AddDiscrete(new IntRange{ Min=0, Max=theWorld.getNumStates() - 1 });
            //Specify that there will be an integer action [0,3]
            RLTypesSpec acts = new RLTypesSpec();
            acts.AddDiscrete(new IntRange{ Cardinality=1, Min=0, Max=3 }); //Cardinality=1 is also default
            //Specify the reward range [-100,10]
            DoubleRange rew = new DoubleRange{ Min=-100.0, Max=10 };
            string extra = "SampleMinesEnvironment(.Net) based on SampleMinesEnvironment(Java) by Brian Tanner.";

            ITaskSpec theTaskSpecObject = TaskSpecFactory.GetTaskSpec("episodic", 1.0, obs, acts, rew, extra);

            return theTaskSpecObject;

        }

        /// <summary>
        /// "Put the environment in a random state and return the appropriate observation." (BTanner)
        /// </summary>
        /// <returns></returns>
        public GeneralState StartEnvironment() {
            if (fixedStartState) {
                bool stateIsValid = theWorld.SetAgentState(startRow, startCol);
                if (!stateIsValid) {
                    theWorld.setRandomAgentState();
                }
            }
            else {
                theWorld.setRandomAgentState();
            }
            GeneralState theObservation = new GeneralState(theWorld.getState());
            return theObservation;
        }

        /// <summary>
        /// "Make sure the action is in the appropriate range, update the state,
        /// generate the new observation, reward, and whether the episode is over." (BTanner)
        /// </summary>
        /// <param name="action"></param>
        /// <param name="outReward"></param>
        /// <param name="outObservation"></param>
        /// <returns></returns>
        public bool StepEnvironment(GeneralAction action, out double outReward, out GeneralState outObservation) {
            /* Make sure the action is valid */
            Debug.Assert(action.IntCount == 1, "Expecting a 1-dimensional integer action. " + action.IntCount + "D was provided");
            Debug.Assert(action.IntValues[0] >= 0, "Action should be in [0,3], " + action.IntValues[0] + " was provided");
            Debug.Assert(action.IntValues[0] < 4, "Action should be in [0,3], " + action.IntValues[0] + " was provided");

            theWorld.UpdatePosition(action.IntValues[0]);

            outObservation = new GeneralState(theWorld.getState());
            bool terminal = theWorld.isTerminal();
            outReward = theWorld.GetReward();

            return terminal;
        }

        /// <summary>
        /// 
        /// </summary>
        public void CleanupEnvironment() {
            /*nothing*/
        }

        /// <summary>
        /// 
        /// </summary>
        /// <param name="message"></param>
        /// <returns></returns>
        public string EnvMessage(string message) {
            /*	Message Description
             * 'set-random-start-state'
             * Action: Set flag to do random starting states (the default)
             */
            if (message.StartsWith("set-random-start-state")) {
                fixedStartState = false;
                return "Message understood.  Using random start state.";
            }

            /*	Message Description
                * 'set-start-state X Y'
                * Action: Set flag to do fixed starting states (row=X, col=Y)
                */
            if (message.StartsWith("set-start-state")) {
                String[] theTokens = message.Split(' ');
                startRow = Int32.Parse(theTokens[1]);
                startCol = int.Parse(theTokens[2]);
                fixedStartState = true;
                return "Message understood.  Using fixed start state.";
            }

            /*	Message Description
            *	'print-state'
            *	Action: Print the map and the current agent location
            */
            if (message.StartsWith("print-state")) {
                theWorld.PrintState();
                return "Message understood.  Printed the state.";
            }

            return "SamplesMinesEnvironment(Java) does not understand your message.";
        }

        #endregion
    }



    /**
     * This class holds all of the internal state information about the environment,
     * and manages the dynamics, state update, reward calculation, etc.
     * @author btanner
     */
    class WorldDescription {

        private readonly int numRows;
        private readonly int numCols;
        public int agentRow;
        public int agentCol;
        private readonly int[,] theMap;
        private Random randGen = new Random();

        public WorldDescription(int[,] worldMap) {
            this.theMap = worldMap;
            
            this.numRows = theMap.Rank;
            this.numCols = theMap.GetLength(1);
        }

        public int getNumStates() {
            return theMap.Length;
        }

        /**
         * Puts the agent into a random state.  Uses a generate and test method, in
         * a loop, only accepts the state if it is valid and not terminal.
         */
        public void setRandomAgentState() {

            int startRow = randGen.Next(numRows);
            int startCol = randGen.Next(numCols);

            while (isTerminal(startRow, startCol) || !isValid(startRow, startCol)) {
                startRow = randGen.Next(numRows);
                startCol = randGen.Next(numCols);
            }

            this.agentRow = startRow;
            this.agentCol = startCol;
        }

        /**
         * Convert the row/col state into a single number.
         * @return
         */
        public int getState() {
            return agentCol * numRows + agentRow;
        }

        /**
         * Sets the agent current state to startRow,startCol.
         * @param startRow
         * @param startCol
         * @return true if the state is valid and not terminal, otherwise
         * return false.
         */
        internal bool SetAgentState(int startRow, int startCol) {
            this.agentRow = startRow;
            this.agentCol = startCol;

            return isValid(startRow, startCol) && !isTerminal();
        }

        public bool isTerminal() {
            return isTerminal(agentRow, agentCol);
        }

        private bool isTerminal(int row, int col) {
            if (theMap[row,col] == SampleMinesEnvironment.WORLD_GOAL || theMap[row,col] == SampleMinesEnvironment.WORLD_MINE) {
                return true;
            }
            return false;
        }

        private bool isValid(int row, int col) {
            bool valid = false;
            if (row < numRows && row >= 0 && col < numCols && col >= 0) {
                if (theMap[row,col] != SampleMinesEnvironment.WORLD_OBSTACLE) {
                    valid = true;
                }
            }
            return valid;
        }

        /**
         * Calculate the reward for the current agent state.
         * @return
         */
        public double GetReward() {
            if (theMap[agentRow,agentCol] == SampleMinesEnvironment.WORLD_GOAL) {
                return 10.0f;
            }

            if (theMap[agentRow,agentCol] == SampleMinesEnvironment.WORLD_MINE) {
                return -100.0f;
            }

            return -1.0f;
        }

        public void UpdatePosition(int theAction) {
            /* When the move would result in hitting an obstacles, the agent simply doesn't move */
            int newRow = agentRow;
            int newCol = agentCol;


            if (theAction == 0) {/*move down*/
                newCol = agentCol - 1;
            }
            if (theAction == 1) { /*move up*/
                newCol = agentCol + 1;
            }
            if (theAction == 2) {/*move left*/
                newRow = agentRow - 1;
            }
            if (theAction == 3) {/*move right*/
                newRow = agentRow + 1;
            }


            /*Check if new position is out of bounds or inside an obstacle */
            if (isValid(newRow, newCol)) {
                agentRow = newRow;
                agentCol = newCol;
            }
        }

        /**
         * Print out the current state to the screen
         */
        internal void PrintState() {
            Console.WriteLine("Agent is at: {0},{1}",agentRow, agentCol);
            Console.WriteLine("Columns:0-10                10-17");
            Console.Write("Col    ");
            for (int col = 0; col < 18; col++) {
                Console.Write("{0} ", col % 10);
            }

            for (int row = 0; row < 6; row++) {
                Console.Write("\nRow: {0} ", row);

                for (int col = 0; col < 18; col++) {
                    if (agentRow == row && agentCol == col) {
                        Console.Write("A ");
                    } else {
                        if (theMap[row,col] == SampleMinesEnvironment.WORLD_GOAL) {
                            Console.Write("G ");
                        }
                        if (theMap[row,col] == SampleMinesEnvironment.WORLD_MINE) {
                            Console.Write("M ");
                        }
                        if (theMap[row,col] == SampleMinesEnvironment.WORLD_OBSTACLE) {
                            Console.Write("* ");
                        }
                        if (theMap[row,col] == SampleMinesEnvironment.WORLD_FREE) {
                            Console.Write("  ");
                        }
                    }
                }
            }
            Console.WriteLine();
        }
    }
}
